package edu.berkeley.nlp.assignments;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

import edu.berkeley.nlp.langmodel.LanguageModel;
import edu.berkeley.nlp.util.Counter;
import edu.berkeley.nlp.util.CounterMap;

public class MyLanguageModel implements LanguageModel {

	static final String STOP = "</S>";
	static final String START = "<S>";

	double total = 0.0;
	// for unigram, c(w')
	Counter<String> wordCounter = new Counter<String>();

	// double totalSizeRev = 0.0;

	double totalBigram = 0.0;
	// for bigram, c*(w,w')
	CounterMap<String, String> bigramCounterMap = new CounterMap<String, String>();
	// CounterMap<String, String> bigramProb = new CounterMap<String, String>();
	Counter<String> contextMap = new Counter<String>();

	// for trigram key:w-2,w-1, value: w
	double totalTrigram = 0.0;
	CounterMap<String, String> trigramCounterMap = new CounterMap<String, String>();
//	Counter<String> preTrigramCounter = new Counter<String>();
	Counter<String> bigramContextMap = new Counter<String>();

	double D = 0.75;

	public static int TYPE_KATZ = 0; // backoff
	public static int TYPE_KN = 1; // backoff
	public static int TYPE_KN_INTERPOLATION = 2; // interpolation
	public static int TYPE_PURE_INTERPOLATION = 3;
	public static int TYPE_PURE_INTERPOLATION_CONTEXT = 4;

	int ngram;
	int type;
	double[] lamda;
	boolean interpolation;

	int UNKPrediction = 0, biUNKPrediction; // n0
	Counter<String> alphaFactor = new Counter<String>();

	public void predictUNKs() {
		for (String word : wordCounter.keySet())
			if (wordCounter.getCount(word) == 1)
				UNKPrediction++;

		if (ngram == 3) {
			for (String wConcat : trigramCounterMap.keySet())
				if (trigramCounterMap.getCounter(wConcat).totalCount() == 1)
					biUNKPrediction++;
		}
	}

	String generateWord() {
		double sample = Math.random();
		double sum = 0.0;
		for (String word : wordCounter.keySet()) {
			sum += getWordProbability(word);
			if (sum > sample) {
				return word;
			}
		}
		return "*UNKNOWN*";
	}

	String generateWord(String wordP) {
		double sample = Math.random();
		double sum = 0.0;
		for (String word : wordCounter.keySet()) {
			sum += getWordProbability(wordP, word);
			if (sum > sample) {
				return word;
			}
		}
		return "*UNKNOWN*";
	}

	String generateWord(String wordPP, String wordP) {
		double sample = Math.random();
		double sum = 0.0;
//		System.out.println("GEN to" + sample);
		for (String word : wordCounter.keySet()) {
			sum += getWordProbability(wordPP, wordP, word);
			// System.out.println(" " + sum);
			// sum += trigramProb.getCount(TrigramKey.getConcatString(wordPP,
			// wordP), word);
			if (sum > sample) {
//				System.out.println(word);
				return word;
			}
		}
		return "*UNKNOWN*";
	}

	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
		String word = generateWord(), lastWord = null, lastlastWord = null;
		// String word = START, lastWord = null, lastlastWord = null;
		if (ngram == 2) {
			while (!word.equals(STOP)) {
				sentence.add(word);
				word = generateWord(word);
			}
		} else if (ngram == 3) {
			if (!word.equals(STOP)) {
				sentence.add(word);
				lastlastWord = word;
				word = generateWord(word);
				lastWord = word;
//				System.out.println("GEN " + lastlastWord + " " + lastWord);
				while (!word.equals(STOP)) {
					sentence.add(word);
					word = generateWord(lastlastWord, lastWord);

					lastlastWord = lastWord;
					lastWord = word;
				}
			}
		}
		return sentence;
	}

	public double getWordProbability(String word) {
		double count = wordCounter.getCount(word);
		if (count == 0) {
			// System.out.println("UNKNOWN WORD: "+sentence.get(index));

			return UNKPrediction / (total + 1.0);
//			return 1.0 / (total + 1.0);
		}
		return count / (total + 1.0);
	}

	public double getWordProbability(String lastWord, String word) {
		double count;
		if (ngram == 3 && type == TYPE_KN) {
			double prob =  getBigramPart(lastWord, word);
			if (prob == 0)
				return getWordProbability(word);
			else 
				return prob;
		}
		
		if (type == TYPE_KN) {
			if (wordCounter.getCount(lastWord) == 0) {
				// UNK
/*	
				if (contextMap.getCount(word) == 0)
					return  1.0
							/ (contextMap.totalCount()+ 1);
				else
					return contextMap.getCount(word)
							/ (contextMap.totalCount() + 1);
	*/
				return getWordProbability(word);
	//			return UNKPrediction / (total + 1.0);
			} else if ((count = bigramCounterMap.getCount(lastWord, word)) > 0) {
				return (count - D) / wordCounter.getCount(lastWord);
			} else {
				double totalDiscount = D
						* bigramCounterMap.getCounter(lastWord).keySet().size()
						/ wordCounter.getCount(lastWord)
						* alphaFactor.getCount(lastWord);

				/*
				 * if (bigramCounterMap.getCounter(lastWord).keySet().size() ==
				 * wordCounter.keySet().size()) totalDiscount = 1; else ///
				 * ????? // totalDiscount = totalDiscount *
				 * alphaFactor.getCount(lastWord); totalDiscount = totalDiscount *
				 * wordCounter.keySet().size() / (wordCounter.keySet().size() -
				 * bigramCounterMap.getCounter(lastWord).keySet().size()) ;
				 */
				double contextCount, contextUnigramProb;

				if ((contextCount = contextMap.getCount(word)) == 0)
					contextUnigramProb = 1 / (contextMap.totalCount() + 1);
				else
					contextUnigramProb = contextCount
							/ (contextMap.totalCount() + 1);

				return totalDiscount * contextUnigramProb;
			}
		} else if (type == TYPE_PURE_INTERPOLATION) {
			double probability = 0.0;

			probability += lamda[0] * getWordProbability(word);

			if (wordCounter.getCount(lastWord) != 0.0)
				probability += (1 - lamda[0])
						* bigramCounterMap.getCount(lastWord, word)
						/ wordCounter.getCount(lastWord);

			// if (probability == 0.0)
			// return 1.0;
			// else
			return probability;
		} else if (type == TYPE_PURE_INTERPOLATION_CONTEXT) {
			double probability = 0.0;

			probability += lamda[0] * getWordProbability(word);


			probability += lamda[1] *getBigramContextPart(lastWord, word);
			
			if (wordCounter.getCount(lastWord) != 0.0)
				probability += (1 - lamda[0] - lamda[1])
						* bigramCounterMap.getCount(lastWord, word)
						/ wordCounter.getCount(lastWord);
			
			// if (probability == 0.0)
			// return 1.0;
			// else
			return probability;
		} else
			return 1.0;
	}

	public double getWordProbability(String lastLastWord, String lastWord,
			String word) {
		double count;
		if (type == TYPE_PURE_INTERPOLATION) {
			double probability = 0.0;

			////
			//// unigram
			probability += lamda[0] * getWordProbability(word);

			////
			//// bigram
			
			probability += lamda[1] * getBigramPart(lastWord, word);

			////
			//// trigram

			probability += (1.0 - lamda[0] - lamda[1]) *getTrigramPart(lastLastWord, lastWord, word);

			return probability;
		} else if (type == TYPE_KN_INTERPOLATION) {
			double probability = 0.0, totalTemp;
			////
			//// discount KN
			
			
			probability += lamda[0] * getKNTrigramDiscountPart(lastLastWord, lastWord, word);

			// 0.095
//			probability += (1-lamda[0]) * getBigramPart(lastWord, word);
			// 0.106
//			probability += (1-lamda[0]) * getBigramContextPart(lastWord, word);

			
			//probability += lamda[1] * getBigramPart(lastWord, word);
			// unigram	
			//probability += (1.0 - lamda[1] - lamda[0]) * getWordProbability(word);
			

			
			probability += (lamda[1]) * getKNBigramDiscountPart(lastWord, word);//getBigramPart(lastWord, word);//getBigramContextPart(lastWord, word);
			// unigram
			probability += (1.0 - lamda[1] - lamda[0]) * getWordProbability(word);
			
			if (probability == 0.0)
				probability = getWordProbability(word);

			
			return probability;
			

		} else if (type == TYPE_KN)  {
			double probability = 0.0;
			probability += getKNTrigramDiscountPart(lastLastWord, lastWord, word);
			
			probability += getKNBrigramContextPart(lastLastWord, lastWord, word);

			if (probability == 0.0)
				probability = getWordProbability(word);

			if (probability == 0.0)
				System.out.println("BAD!!" + probability);
			
			return probability;
		}else 
			return 1.0;
	}
	
	/* 
	 * KN 3
	 */
	public double getKNTrigramDiscountPart(String lastLastWord, String lastWord, String word ) {
		double probability = 0.0, totalTemp;
		if ((totalTemp = trigramCounterMap.getCounter(TrigramKey
				.getConcatString(lastLastWord, lastWord)).totalCount()) != 0.0) {
//			if (trigramCounterMap.getCount(TrigramKey.getConcatString(
//					lastLastWord, lastWord), word) != 0.0)
				probability = (Math.max(trigramCounterMap.getCount(
						TrigramKey.getConcatString(lastLastWord, lastWord),
						word)
						- D, 0.0))
						/ (totalTemp + 1);
//			else if (!wordCounter.containsKey(word))
//				probability = getWordProbability(word);
//			else
//				probability = 1.0 / (totalTemp + 1);

//				if (probability == 0.0)
//					System.out.println("BAD1!!");
		} else {
			///
			/// with this, a little better for 3-correlation-bigram, not for 3-correlation-trigram
			///
//			probability = getWordProbability(word);
		}
//		if (probability == 0.0)
//			System.out.println("BAD1-2!!");
		return probability;
	}

	public double getBigramContextPart(String lastWord, String word) {
		double probability = 0.0, totalTemp, temp;
			// turn to bigram context map
			if ((temp = bigramContextMap.getCount(TrigramKey
					.getConcatString(lastWord, word))) == 0) {
//				probability = 1.0 / (bigramContextMap.totalCount() + 1.0);
			} else {
				probability = temp / (bigramContextMap.totalCount() + 1.0);
			}
			
			
		
		return probability;
	}
	
	public double getUnigramContextPart(String word) {
		double contextCount = 0.0;
		if ((contextCount = contextMap.getCount(word)) == 0)
			return 1 / (contextMap.totalCount() + 1);
		else
			return contextCount
					/ (contextMap.totalCount() + 1);	
	}
	
	public double getKNBrigramContextPart(String lastLastWord, String lastWord, String word ) {
		double probability = 0.0, totalTemp, temp;
		if ((totalTemp = trigramCounterMap.getCounter(
				TrigramKey.getConcatString(lastLastWord, lastWord))
				.totalCount()) != 0.0) {
			// turn to bigram context map
			if ((temp = bigramContextMap.getCount(TrigramKey
					.getConcatString(lastWord, word))) == 0) {
//				probability = 1.0 / (bigramContextMap.totalCount() + 1.0);
			} else {
				probability = temp / (bigramContextMap.totalCount() + 1.0);
			}
			
			probability *= D * trigramCounterMap.getCounter((TrigramKey
					.getConcatString(lastLastWord, lastWord))).keySet().size() / totalTemp;

		} else {
			probability = getBigramPart(lastWord, word);

		}

//		if (probability == 0.0)
//			probability = getWordProbability(word);

//		if (probability == 0.0)
//			System.out.println("BAD2!!");
		
		return probability;
	}
	
	/*
	 * 
	 */
	
	
	
	
	public double getTrigramPart(String lastLastWord, String lastWord, String word) {
		double probability = 0.0;

		double totalTemp;
		if ((totalTemp = trigramCounterMap.getCounter(
				TrigramKey.getConcatString(lastLastWord, lastWord))
				.totalCount()) != 0.0) {			
				probability = trigramCounterMap.getCount(TrigramKey
							.getConcatString(lastLastWord, lastWord), word)
					/ totalTemp;

			
		} else {
			//probability = getBigramPart(lastWord, word);
			
			// turn to bigram context map
			/*
			if ((totalTemp = bigramContextMap.getCount(TrigramKey.getConcatString(
					lastWord, word))) == 0) {
//				probability = 1.0
//						/ (bigramContextMap.totalCount() + 1.0);
			} else {
				probability = totalTemp / (bigramContextMap.totalCount() + 1.0);
			}*/
		}
		return probability;
	}


	public double getUnigramPart(String word) {
		double count = wordCounter.getCount(word);
		if (count == 0) {
			// System.out.println("UNKNOWN WORD: "+sentence.get(index));

			return UNKPrediction - D/ (total + 1.0);
//			return 1.0 / (total + 1.0);
		}
		return count - D / (total + 1.0);
	}
	public double getBigramPart(String lastWord, String word) {
		double probability = 0.0;
		if (wordCounter.getCount(lastWord) != 0.0) {
			if (bigramCounterMap.getCount(lastWord, word) == 0) {
				//// addictive smoothing is so wrong!!! awful1!!
				/*probability = 1.0
					/ (wordCounter.getCount(lastWord) + 1.0);*/
			} else 
				probability = 
					bigramCounterMap.getCount(lastWord, word)
				/ (wordCounter.getCount(lastWord) + 1.0);
	
		} else {
			//probability = getWordProbability( word);

			
			if (contextMap.getCount(word) == 0)
				probability = 1.0
						/ (contextMap.totalCount()+ 1);
			else
				probability = contextMap.getCount(word)
						/ (contextMap.totalCount() + 1);
			
				// UNK ???
//				probability += 1.0/ (total + 1.0);
			//	probability = UNKPrediction/ (total + 1.0);
			
		}
		return probability;
	}
	public double getKNBigramDiscountPart(String lastWord, String word) {
		double probability = 0.0;
		if (wordCounter.getCount(lastWord) != 0.0) {
			if (bigramCounterMap.getCount(lastWord, word) == 0) {
				//// addictive smoothing is so wrong!!! awful1!!
				/*probability = 1.0
					/ (wordCounter.getCount(lastWord) + 1.0);*/
			} else 
/*				probability = 
					bigramCounterMap.getCount(lastWord, word)
				/ (wordCounter.getCount(lastWord) + 1.0);
	*/		
				probability = 
					Math.max(bigramCounterMap.getCount(lastWord, word) - D, 0)
				/ (wordCounter.getCount(lastWord) + 1.0);
		} else {
			//probability = getWordProbability( word);

			
			if (contextMap.getCount(word) == 0)
				probability = 1.0
						/ (contextMap.totalCount()+ 1);
			else
				probability = contextMap.getCount(word)
						/ (contextMap.totalCount() + 1);
			
				// UNK ???
//				probability += 1.0/ (total + 1.0);
			//	probability = UNKPrediction/ (total + 1.0);
			
		}
		return probability;
	}
	
	public double getWordProbability(List<String> sentence, int index) {
		String word = sentence.get(index);
		return getWordProbability(word);

	}

	public double getBigramProbability(List<String> sentence, int index) {
		String wordP = sentence.get(index - 1), word = sentence.get(index);

		return getWordProbability(wordP, word);

	}

	// TODO
	public double getTrigramProbability(List<String> sentence, int index) {
		String wordPP = sentence.get(index - 2), wordP = sentence
				.get(index - 1), word = sentence.get(index);

		return getWordProbability(wordPP, wordP, word);
	}

	public void setLamda(double[] lamda) {
		this.lamda = new double[lamda.length];
		for (int i = 0; i < lamda.length; i++)
			this.lamda[i] = lamda[i];
	}

	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(STOP);
		double probability = 1.0;
		if (ngram == 2) {
			probability *= getWordProbability(stoppedSentence, 0);
			for (int index = 1; index < stoppedSentence.size(); index++) {
				probability *= getBigramProbability(stoppedSentence, index);
			}
		} else if (ngram == 3) {
			probability *= getWordProbability(stoppedSentence, 0);
			probability *= getBigramProbability(stoppedSentence, 1);
			for (int index = 2; index < stoppedSentence.size(); index++) {
				probability *= getTrigramProbability(stoppedSentence, index);
			}
		}
		return probability;
	}

	public void KNBigramBackoff() {

		// alphaFactor

		for (String wP : bigramCounterMap.keySet()) {
			double zeroPortion = 0, nonZeroPortion = 0;
			for (String w : bigramCounterMap.getCounter(wP).keySet()) {
				if (bigramCounterMap.getCount(wP, w) == 0)
					zeroPortion += contextMap.getCount(w);
				else
					nonZeroPortion += contextMap.getCount(w);
			}
			if (zeroPortion == 0)
				alphaFactor.incrementCount(wP, 1);
			else
				alphaFactor.incrementCount(wP,
						(zeroPortion + nonZeroPortion + 1) / zeroPortion);
		}

		// double discountAmount, count, prob;
		/*
		 * for (String wP : bigramCounterMap.keySet()) { double totalTemp =
		 * bigramCounterMap.getCounter(wP).totalCount(); // discount phase //
		 * discountAmount = 0; // for (String w : wordCounter.keySet()) { for
		 * (String w : bigramCounterMap.getCounter(wP).keySet()) { totalDiscount +=
		 * 0.75 / totalTemp; if ((count = bigramCounterMap.getCount(wP, w)) > 0) {
		 * prob = Math.max(count - D, 0) / totalTemp; discountAmount += (count -
		 * Math.max(count - D, 0)) / totalTemp; // System.out.print("original
		 * prob " + count / totalTemp); // bigramProb.setCount(wP, w, prob); //
		 * System.out.println(" discounted prob " + prob); } } // backoff phase }
		 */
		// System.out.println("fininish preparation phase");
		/*
		 * // backoff phase double disbributionDenominator = 0; for (String w :
		 * wordCounter.keySet()) { disbributionDenominator +=
		 * revBigramCounterMap.getCounter(w) .keySet().size(); } for (String wP :
		 * wordCounter.keySet()) { for (String w : wordCounter.keySet()) {
		 * double count; if ((count = bigramCounterMap.getCount(wP, w)) == 0) {
		 * double prob = 0.0; prob = discountAmount
		 * revBigramCounterMap.getCounter(w).keySet().size() /
		 * disbributionDenominator; bigramProb.setCount(wP, w, prob); } } }
		 */
	}

	public void KNTrigramStat() {
		/*
		// denominator in prob calculation
		for (String wConcat : trigramCounterMap.keySet()) {
			for (String w : trigramCounterMap.getCounter(wConcat).keySet()) {
				preTrigramCounter.incrementCount(wConcat, trigramCounterMap
						.getCount(wConcat, w));
			}
		}*/
	}

	public MyLanguageModel(Collection<List<String>> sentenceCollection,
			int ngram, int type, boolean useLeftOutSet, int[] lamda) {
		this.ngram = ngram;
		this.type = type;
		this.interpolation = useLeftOutSet;
		String lastWord, lastLastWord;
		for (List<String> sentence : sentenceCollection) {
			List<String> stoppedSentence = new ArrayList<String>(sentence);
			// stoppedSentence.add(0, START);
			stoppedSentence.add(STOP);
			lastWord = "";
			lastLastWord = "";

			for (int index = 0; index < stoppedSentence.size(); index++) {
				String word = stoppedSentence.get(index);
				if (index > 0)
					lastWord = stoppedSentence.get(index - 1);
				if (index > 1)
					lastLastWord = stoppedSentence.get(index - 2);
//			for (String word : stoppedSentence) {
				wordCounter.incrementCount(word, 1.0);

				if (ngram >= 2 && !lastWord.equals("")) {
					bigramCounterMap.incrementCount(lastWord, word, 1.0);
					if (bigramCounterMap.getCount(lastWord, word) == 1.0)
						contextMap.incrementCount(word, 1.0);
				}

				if (ngram == 3 && !lastLastWord.equals("")) {
					trigramCounterMap
							.incrementCount(TrigramKey.getConcatString(
									lastLastWord, lastWord), word, 1.0);
					if (trigramCounterMap.getCount(TrigramKey.getConcatString(
							lastLastWord, lastWord), word) == 1.0)
						bigramContextMap.incrementCount(TrigramKey
								.getConcatString(lastWord, word), 1.0);
				}
//				lastLastWord = lastWord;
//				lastWord = word;
			}
		}
		total = wordCounter.totalCount();
		totalBigram = bigramCounterMap.totalCount();
		// totalTrigram = trigramCounterMap.totalCount();

		predictUNKs();

		if (type == TYPE_KATZ) {

		} else if (type == TYPE_KN) {

			if (ngram == 2) {
				KNBigramBackoff();
			}
		} else if (type == TYPE_KN_INTERPOLATION) {
			if (ngram == 3)
				KNTrigramStat();
		} else if (type == TYPE_PURE_INTERPOLATION) {

		}
		// if (ngram == 3)
		// KneserNeyTrigramSmoothing();

		System.out.println("finish ngram counter");
	}

}

class TrigramKey {
	// trigram key manipulation
	static String CONCAT = "<concat>";

	public static String getFirstString(String concatString) {
		for (int i = 0; i < concatString.length(); i++) {
			if (concatString.substring(0, i).endsWith(CONCAT))
				return concatString.substring(0, i - 1 - CONCAT.length());
		}
		return null;
	}

	public static String getSecondString(String concatString) {
		for (int i = 0; i < concatString.length(); i++) {
			if (concatString.substring(0, i).endsWith(CONCAT))
				return concatString.substring(i);
		}
		return null;
	}

	public static String getConcatString(String first, String second) {
		String concat = "";
		concat = first + CONCAT + second;
		return concat;
	}

}